from datetime import time, datetime

from db_helper.milvus_database_helper import MilvusDatabase
from sentence_transformers import SentenceTransformer
from pdf_analyse.pdf_fitz_analyse import PdfFitzAnalyse
import re
import threading
import sys

# pdf相关信息
pdf_path_sea_ship = '../pdf/钢质海船入级规范2022.pdf'
file_name = "钢质海船入级规范2022.pdf"
# 用于标题中的页码标记
db = MilvusDatabase(dim=512, host="117.78.5.132", port="19530", user="root", password="Milvus@aI", database="test")
# 获取所有章节及其内容
analyse = PdfFitzAnalyse(pdf_path_sea_ship)
chapters_result = analyse.get_chapters_and_content()

result_dict = []
classpath = sys.path
sentence_transformer = SentenceTransformer(
    model_name_or_path=f'{classpath[0o4]}/sentence-transformers/distiluse-base-multilingual-cased-v1',
    tokenizer_kwargs={"clean_up_tokenization_spaces": False})


def insert_line(_title, _content, _db, _sentence_transformer):
    current_page_num = 0
    title_match = re.search(f"{analyse.get_page_num_prefix()}(\\d+)", _title)
    if title_match:
        current_page_num = int(title_match.group(1))
    line_list = [item for item in _content.split("\n") if not len(item.replace(" ", "")) == 0]
    embeddings = _sentence_transformer.encode(line_list)
    data_list = [{"file_name": file_name, "page": current_page_num, "content_code": f'{_title}_line_{i}',
                  "embeddings": embeddings[i], "content": line_list[i]} for i in range(len(embeddings))]
    db.insert_data(entities=data_list)


threads = []
start_time = datetime.now()
# 记录入库
for title, content in chapters_result.items():
    thread = threading.Thread(target=insert_line, args=(title, content, db, sentence_transformer))
    thread.start()
    threads.append(thread)
    thread.join()
end_time = datetime.now()
print(end_time - start_time)



